package com.zoe.onelink.logging.trace.sql;

import brave.Span;
import brave.Tracer;
import cn.hutool.core.util.NumberUtil;
import com.alibaba.fastjson.JSON;
import com.zoe.onelink.common.constant.OnelinkMetricsConstant;
import com.zoe.onelink.common.constant.TraceSpanConstant;
import com.zoe.onelink.logging.trace.property.TraceLogProperties;
import com.zoe.onelink.logging.trace.util.SqlParserUtil;
import com.zoe.onelink.logging.trace.util.TraceLogUtil;
import com.zoe.onelink.util.PathMatchUtil;
import com.zoe.onelink.util.StrUtil;
import io.micrometer.core.instrument.Metrics;
import io.micrometer.core.instrument.Tags;
import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.lang.reflect.InvocationTargetException;
import java.util.Collection;
import java.util.Iterator;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * <p>标题: </p>
 * <p>描述: </p>
 * <p>版权: Copyright (c) 2021</p>
 * <p>公司: 智业软件股份有限公司</p>
 *
 * @version: 1.0
 * @author: ruansheng
 * @date: 2021-08-09
 */
@Intercepts({
        @Signature(
                type = Executor.class,
                method = "query",
                args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}),
        @Signature(
                type = Executor.class,
                method = "query",
                args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class}),
        @Signature(
                type = Executor.class,
                method = "queryCursor",
                args = {MappedStatement.class, Object.class, RowBounds.class}),
        @Signature(
                type = Executor.class,
                method = "update",
                args = {MappedStatement.class, Object.class})
})
@Slf4j
@RequiredArgsConstructor
public class SqlTraceLogInterceptor implements Interceptor {

    private final TraceLogProperties.SqlLogProperties sqlLogProperties;
    private final Tracer tracer;
    private final Map<String, SqlTuple> sqlHashMap = new ConcurrentHashMap<>();
    private static final String SQL_COUNT_ID_SUFFIX = "_mpCount";

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        Object[] args = invocation.getArgs();
        // 获得SQL语句
        MappedStatement mappedStatement = (MappedStatement) args[0];
        Span span = createNextTraceSpan(mappedStatement);
        String returnRows = "-1";
        try {
            Object proceed = invocation.proceed();
            returnRows = this.getReturnRows(proceed);
            String sqlResult = this.getCountSqlResult(mappedStatement, proceed);
            span.tag(TraceSpanConstant.SQL_ROWS, returnRows);
            span.tag(TraceSpanConstant.SQL_RESULT, sqlResult);
            return proceed;
        } catch (Throwable ex) {
            Throwable cause = ex.getCause();
            if (ex instanceof InvocationTargetException && cause != null) {
                span.error(cause);
            } else {
                span.error(ex);
            }
            throw ex;
        } finally {
            this.logSql(args, mappedStatement, returnRows, span);
            SqlExecuteContextHolder.remove();
            span.finish();
        }
    }

    /**
     * 获取COUNT语法返回结果
     *
     * @param mappedStatement
     * @param result
     * @return
     */
    private String getCountSqlResult(MappedStatement mappedStatement, Object result) {
        if (!StrUtil.endWith(mappedStatement.getId(), SQL_COUNT_ID_SUFFIX, true)) {
            return "";
        }
        String sqlResult;
        if (result instanceof Collection) {
            Iterator<?> it = ((Collection<?>) result).iterator();
            if (it.hasNext()) {
                Object retVal = it.next();
                sqlResult = retVal instanceof Number ? retVal.toString() : "";
            } else {
                sqlResult = "";
            }
        } else if (result instanceof Number) {
            sqlResult = result.toString();
        } else {
            sqlResult = "";
        }
        return sqlResult;
    }

    /**
     * 获取SQL返回行数
     */
    private String getReturnRows(Object proceed) {
        return SqlExecuteContextHolder.get(SqlExecuteContextHolder.ROWS).orElseGet(() -> {
            String size;
            if (proceed instanceof Collection) {
                size = ((Collection<?>) proceed).size() + "";
            } else if (proceed instanceof Number) {
                size = proceed + "";
            } else if (proceed instanceof String) {
                size = NumberUtil.isNumber(proceed.toString()) ? NumberUtil.parseLong(proceed.toString()) + "" : "";
            } else {
                size = "";
            }
            return size;
        });
    }

    private void logSql(Object[] args, MappedStatement mappedStatement, String returnRows, Span span) {
        try {
            // 方法路径匹配
            String sqlId = mappedStatement.getId();
            // 语法是否在白名单中
            boolean isPermit = PathMatchUtil.isPermit(sqlId, this.sqlLogProperties.getIncludeMethods(),
                    this.sqlLogProperties.getExcludeMethods(), true, true, "SQL_TRACE");
            // 黑名单语法直接忽略
            if (!isPermit) {
                if (log.isDebugEnabled()) {
                    log.debug("SQL处于日志黑名单中,已自动忽略. SQL_ID:{}", sqlId);
                }
            }
            // SQL语法、SQL参数
            String sqlGrammar, sqlParam;
            // 是否是Count语法
            boolean isCountSql = StrUtil.endWith(sqlId, SQL_COUNT_ID_SUFFIX, true);
            // SQL语法缓存
            Optional<String> cachedSql = SqlExecuteContextHolder.get(SqlExecuteContextHolder.SQL);
            // COUNT语法默认不记录
            if (isCountSql) {
                sqlGrammar = "";
                sqlParam = "";
                // SQL语法缓存存在则直接使用缓存
            } else if (cachedSql.isPresent()) {
                sqlGrammar = cachedSql.get();
                sqlParam = SqlExecuteContextHolder.get(SqlExecuteContextHolder.PARAM).orElse("");
                // 记录可直接执行的SQL语法
            } else if (TraceLogProperties.SqlLogProperties.Strategy.FILLED.equals(this.sqlLogProperties.getStrategy())) {
                BoundSql boundSql = this.getBoundSql(args, mappedStatement);
                sqlGrammar = this.getFilledSql(mappedStatement, boundSql);
                sqlParam = "";
                // SQL语法与参数分开记录
            } else if (TraceLogProperties.SqlLogProperties.Strategy.PREPARED.equals(this.sqlLogProperties.getStrategy())) {
                BoundSql boundSql = this.getBoundSql(args, mappedStatement);
                sqlGrammar = this.getBoundSql(args, mappedStatement).getSql();
                sqlParam = TraceLogUtil.trim(JSON.toJSONString(boundSql.getParameterObject()), sqlLogProperties.getTrimType(), sqlLogProperties.getParamMaxLength());
                // 只记录预编译语法
            } else if (TraceLogProperties.SqlLogProperties.Strategy.PREPARED_ONLY.equals(this.sqlLogProperties.getStrategy())) {
                sqlGrammar = this.getBoundSql(args, mappedStatement).getSql();
                sqlParam = "";
                // 其他未知情况
            } else {
                sqlGrammar = "unknown";
                sqlParam = "unknown";
            }
            if (log.isDebugEnabled() && !isCountSql) {
                log.debug("SQL返回行数:{} , SQL语法: \n{} \n{}", returnRows, sqlGrammar.replace("\n", ""), sqlParam);
            }
            String sqlGrammarLog = TraceLogUtil.trim(sqlGrammar, sqlLogProperties.getTrimType(), sqlLogProperties.getGrammarMaxLength());
            String sqlParamLog = TraceLogUtil.trim(sqlParam, sqlLogProperties.getTrimType(), sqlLogProperties.getParamMaxLength());
            String consume = SqlExecuteContextHolder.get(SqlExecuteContextHolder.CONSUME).orElse("");
            SqlTuple sqlTuple = this.sqlHashMap.computeIfAbsent(sqlId, key -> {
                String preparedSql = getBoundSql(args, mappedStatement).getSql();
                return new SqlTuple(preparedSql.length(), preparedSql.hashCode());
            });
            span.tag(TraceSpanConstant.SQL_GRAMMAR, sqlGrammarLog);
            span.tag(TraceSpanConstant.SQL_PARAM, sqlParamLog);
            span.tag(TraceSpanConstant.SQL_HASH, String.valueOf(sqlTuple.getHashCode()));
            span.tag(TraceSpanConstant.SQL_CONSUME, consume);
            // 记录预编译语法长度即可,可执行语法长度可以根据整个日志长度估算
            span.tag(TraceSpanConstant.SQL_LENGTH, String.valueOf(sqlTuple.getLength()));
            this.metrics(sqlId, consume);
        } catch (Exception e) {
            span.tag(TraceSpanConstant.SQL_GRAMMAR, "log failure");
        }

    }

    private BoundSql getBoundSql(Object[] args, MappedStatement mappedStatement) {
        return mappedStatement.getBoundSql(args[1]);
    }

    /**
     * 指标
     */
    private void metrics(String sqlId, String consume) {
        long consumeMss = (long) (Double.parseDouble(StrUtil.emptyToDefault(consume, "0")) * 1000);
        Tags tags = Tags.of("sqlId", sqlId);
        Metrics.counter(OnelinkMetricsConstant.Sql.SQL_EXEC_COUNT, tags).increment();
        // 执行时间超过3秒
        if (consumeMss > 3_000_000) {
            Metrics.timer(OnelinkMetricsConstant.Sql.SQL_CONSUME, tags)
                    .record(consumeMss, TimeUnit.MICROSECONDS);
        }
    }

    /**
     * 获取填充过参数的SQL语句
     */
    private String getFilledSql(MappedStatement mappedStatement, BoundSql boundSql) {
        try {
            return SqlParserUtil.fillSqlParam(mappedStatement.getConfiguration(), boundSql);
        } catch (Exception e) {
            log.warn("填充SQL语法参数失败:{}", e.getMessage());
            return boundSql.getSql();
        }
    }

    /**
     * 创建下一个Zipkin链路Span
     */
    private Span createNextTraceSpan(MappedStatement mappedStatement) {
        Span span = this.tracer.nextSpan().kind(Span.Kind.SERVER).start();
        String fullSqlId = mappedStatement.getId();
        String simpleSqlId = TraceLogUtil.abbreviator(fullSqlId);
        String spanName = StrUtil.format("[SQL] [{}] {}", mappedStatement.getSqlCommandType().name(), simpleSqlId);
        span.name(spanName);
        span.tag(TraceSpanConstant.LOG_TYPE, TraceSpanConstant.LogType.SQL);
        span.tag(TraceSpanConstant.SQL_ID, fullSqlId);
        span.tag(TraceSpanConstant.SQL_TYPE, mappedStatement.getSqlCommandType().name());
        return span;
    }

    @Getter
    @AllArgsConstructor
    private static class SqlTuple {

        private int length;
        private int hashCode;

    }


}
